import os
import shutil
import subprocess
import sys


def compile_py_to_pyd(folder_path):
    # 检查 pyarmor 是否已安装
    try:
        subprocess.run(["pyarmor", "--version"], capture_output=True, check=True)
    except (subprocess.SubprocessError, FileNotFoundError):
        print("错误: 未找到 pyarmor 命令。请确保已安装 pyarmor 并添加到系统路径中。")
        print("可以使用命令安装: pip install pyarmor")
        return

    # Step 1: pyarmor cfg restrict_module=0
    setup = os.path.join(folder_path, "setup.py")
    subprocess.run(["pyarmor", "cfg", "restrict_module=0"])

    # 获取传入文件夹路径的上一级目录
    parent_dir = os.path.dirname(folder_path)

    # Step 2: Loop through all .py files in the folder
    for filename in os.listdir(folder_path):
        # 去掉.txt后缀
        date_str = os.path.splitext(filename)[0]
        if date_str == "__init__":
            break
        pyx = os.path.join(folder_path, date_str + ".pyx")
        pyc = os.path.join(folder_path, date_str + ".c")
        if filename.endswith(".py"):
            py_file = os.path.join(folder_path, filename)

            # pyarmor gen <filename>
            subprocess.run(["pyarmor", "gen", py_file])

            # 在上一级目录中找到生成的 dist 文件夹
            dist_dir = os.path.join(parent_dir, "dist")
            py_file_in_dist = os.path.join(dist_dir, filename)
            pyx_file = py_file_in_dist.replace(".py", ".pyx")

            if os.path.exists(py_file_in_dist):
                # Rename the file
                os.rename(py_file_in_dist, pyx_file)

                # Move dist contents to the project root
                for item in os.listdir(dist_dir):
                    src = os.path.join(dist_dir, item)
                    dst = os.path.join(folder_path, item)
                    shutil.move(src, dst)

                # Clean up the dist directory
                shutil.rmtree(dist_dir)

                # Create a setup.py file
                setup_code = f"""
from distutils.core import setup
from Cython.Build import cythonize

setup(
    ext_modules=cythonize(r"{pyx}")
)
"""
                setup_file = os.path.join(folder_path, "setup.py")
                with open(setup_file, "w") as f:
                    f.write(setup_code)

                # Step 3: Run the setup.py to compile to .pyd
                subprocess.run(["python", setup, "build_ext", "--inplace"])

                # Clean up setup.py
                os.remove(setup_file)
                os.remove(pyx)
                os.remove(pyc)
                shutil.rmtree(os.path.join(folder_path, "pyarmor_runtime_000000"))
            else:
                print(f"File not found: {py_file_in_dist}")


# Example usage
dir_list=['air','sea'] # 脚本文件所在目录
folder_path = r"D:\code\vue_rust_py\zhang_xiu_xian-main\zhang_xiuxian\src-tauri\python_pys"
for dir in dir_list: # 循环遍历目录列表
    dir_path=os.path.join(folder_path,dir)
    compile_py_to_pyd(dir_path)
print("ok")